from typing import Annotated

from langchain.chat_models import init_chat_model
from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI

import os

def create_llm() -> ChatOpenAI:
    """创建并返回配置好的大模型实例"""
    return ChatOpenAI(
        api_key=os.getenv("DASHSCOPE_API_KEY") or "your_api_key_here",
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
        model="qwen2.5-math-72b-instruct"
    )


class State(TypedDict):
    messages: Annotated[list, add_messages]


graph_builder = StateGraph(State)


llm = create_llm()


def chatbot(state: State):
    # 这个返回值会更新状态（State）中的 "messages" 字段，把模型的回复添加进去。
    return {"messages": [llm.invoke(state["messages"])]}


# The first argument is the unique node name
# The second argument is the function or object that will be called whenever
# the node is used.
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph = graph_builder.compile()
from IPython.display import Image, display

try:
    display(Image(graph.get_graph().draw_mermaid_png()))
except Exception:
    # This requires some extra dependencies and is optional
    pass


def stream_graph_updates(user_input: str):
    ## graph.stream(...)：LangGraph 提供的方法，用于以“流式”方式运行整个图
    for event in graph.stream({"messages": [{"role": "user", "content": user_input}]}):
        ## for value in event.values() 遍历到的是节点的输出数据。
        for value in event.values():
            ##value["messages"][-1].content 表示取出最新的那条消息的内容（也就是模型的回答）并打印出来。
            print("Assistant:", value["messages"][-1].content)


while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break
        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break